#ifndef GM_SIMPLE_1_CC
#define GM_SIMPLE_1_CC
#include<string>
#include<fstream>
#include<sstream>
#include<iostream>
#include<algorithm>
#include<iterator>
#include<vector>
#include<map>
#include<cstdlib>
#include<ctime>
#include<cmath>
#include "def.hpp"
#include "utils.hpp"
#include "topicmodel_geo_abstract.hpp"
#include "boost/unordered_map.hpp" // for hashtable

typedef map<int, double> UserGeoType;
typedef vector<UserGeoType> UserGeoVector;

class GeoSimple1Batch:public GeoTopicModel{
protected:
  UserGeoVector var_lambda; // the variational parameters for per-user regional distribution
  double** var_nu; // the variatinoal parameters for per-word topic distribution
  double** var_zeta; // the variational parameters for per-topic word distribution
  double* var_zeta_k;
  double** var_xi; // the variational parameters for per-region topic distribution
  double* var_xi_r;
  double* alpha; // the prior distribution over regional topic distribution
  double* beta; // the prior distribution over topic distribution
  double* eta; // the prior distribution over user region distribution
  double AlphaSum;
  double BetaSum;

  // caches
  double** cache_digamma_var_zeta;
  double* cache_digamma_var_zeta_k;
  double** cache_digamma_var_xi;
  double* cache_digamma_var_xi_r;

  // statistics
  double** ss_zeta;
  double** ss_xi;

  int max_doc_len;
  int var_iter;
  double lowerbound;

  ofstream file_output;
public:
  GeoSimple1Batch();
  virtual void Inference();
  virtual void Prediction();
  virtual double Likelihood();
  virtual int ReadDOCS();
  virtual int GetDocuments(int); // this indicate how documents are read through into the model
  virtual int SaveModel();

private:
  int Single_E_Step(int);
  int Single_M_Step();
  void Save_Expectation_Phi_VarZeta();
  void Save_Expectation_Pi_VarXi();
};

int GeoSimple1Batch::SaveModel(){
  string phi_file_name = PuzaDEF::Instance()->output_file_name + ".phi";
  file_output.open(phi_file_name.c_str());
  Save_Expectation_Phi_VarZeta();
  file_output.close();
  string pi_file_name = PuzaDEF::Instance()->output_file_name + ".pi";
  file_output.open(pi_file_name.c_str());
  Save_Expectation_Pi_VarXi();
  file_output.close();

  /*   string pi_file_name = PuzaDEF::Instance()->output_file_name + ".map";
  file_output.open(pi_file_name.c_str());
  for(unsigned int i=0; i < geo_mapping.size(); i++){
    file_output << i << " " << geo_mapping[i] << endl;
  }
  file_output.close(); */

  return 0;
}

GeoSimple1Batch::GeoSimple1Batch(){
  // allocate variational parameters
  var_lambda.clear();
  allocate_memory<double>(var_zeta,PuzaDEF::Instance()->TOPIC_NUM,PuzaDEF::Instance()->TERM_NUM);
  allocate_memory<double>(ss_zeta,PuzaDEF::Instance()->TOPIC_NUM,PuzaDEF::Instance()->TERM_NUM);
  allocate_memory<double>(cache_digamma_var_zeta,PuzaDEF::Instance()->TOPIC_NUM,PuzaDEF::Instance()->TERM_NUM);
  allocate_memory<double>(var_xi,PuzaDEF::Instance()->GEO_REGION_NUM,PuzaDEF::Instance()->TOPIC_NUM);
  allocate_memory<double>(ss_xi,PuzaDEF::Instance()->GEO_REGION_NUM,PuzaDEF::Instance()->TOPIC_NUM);
  allocate_memory<double>(cache_digamma_var_xi,PuzaDEF::Instance()->GEO_REGION_NUM,PuzaDEF::Instance()->TOPIC_NUM);
  // allocate prior parameters
  alpha = new double[PuzaDEF::Instance()->TOPIC_NUM];
  beta = new double[PuzaDEF::Instance()->TERM_NUM];
  eta = new double[PuzaDEF::Instance()->GEO_REGION_NUM];
  var_zeta_k = new double[PuzaDEF::Instance()->TOPIC_NUM];
  var_xi_r = new double[PuzaDEF::Instance()->GEO_REGION_NUM];
  cache_digamma_var_zeta_k = new double[PuzaDEF::Instance()->TOPIC_NUM];
  cache_digamma_var_xi_r = new double[PuzaDEF::Instance()->GEO_REGION_NUM];

  BetaSum = 0.0;
  for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
    beta[v] = PuzaDEF::Instance()->LDA_initial_beta;
    BetaSum = BetaSum + beta[v];
  }

  AlphaSum = 0.0;
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    alpha[k] = PuzaDEF::Instance()->LDA_initial_alpha;
    AlphaSum = AlphaSum + alpha[k];
  }

  // some initialization
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    var_zeta_k[k] = 0.0;
    for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
      var_zeta[k][v] =  random_gamma_distribution(100,0.01);
      var_zeta_k[k] = var_zeta_k[k] + var_zeta[k][v];
      
      cache_digamma_var_zeta[k][v] = digamma(var_zeta[k][v]);
    }
    cache_digamma_var_zeta_k[k] = digamma(var_zeta_k[k]);
  }

  for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    var_xi_r[r] = 0.0;
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      var_xi[r][k] = random_gamma_distribution(100,0.01);
      var_xi_r[r] = var_xi_r[r] + var_xi[r][k];
      
      cache_digamma_var_xi[r][k] = digamma(var_xi[r][k]);
    }
    cache_digamma_var_xi_r[r] = digamma(var_xi_r[r]);
  }
}

int GeoSimple1Batch::Single_E_Step(int d){
  int small_iter = 0;
  TermList current_list = TM_DOCS[d];
  int term_count = current_list.size();

  // E-step for one single document
  double var_nu_sum = 0;
  int current_region = doc_geo_stamp[d];
  for(int i=0; i < term_count; i++){
    int term_id = current_list[i].term_id;
    var_nu_sum = 0;
    // calculate un-normalized nu
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      var_nu[i][k] = cache_digamma_var_zeta[k][term_id] - cache_digamma_var_zeta_k[k] + cache_digamma_var_xi[current_region][k] - cache_digamma_var_xi_r[current_region]; // in log space
      if(k == 0){
          var_nu_sum = var_nu[i][k];
      }
        else{
          var_nu_sum = log_sum(var_nu_sum,var_nu[i][k]);
        }
    }
    // normalize nu
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      var_nu[i][k] = exp(var_nu[i][k] - var_nu_sum);
    }
  }
  
  // gather statistics and update lower bound
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    for(int i=0; i < term_count; i++){
      int term_id = current_list[i].term_id;
      int count = current_list[i].term_stat;
      ss_zeta[k][term_id] = ss_zeta[k][term_id] + count * var_nu[i][k];
      ss_xi[current_region][k] = ss_xi[current_region][k] + count * var_nu[i][k];

      lowerbound = lowerbound + count * var_nu[i][k] * (cache_digamma_var_zeta[k][term_id] - cache_digamma_var_zeta_k[k]);
      lowerbound = lowerbound + count * var_nu[i][k] * (cache_digamma_var_xi[current_region][k] - cache_digamma_var_xi_r[current_region]);      
      lowerbound = lowerbound - count * var_nu[i][k] * log(var_nu[i][k]);
    }
  }

  return small_iter;
}

int GeoSimple1Batch::Single_M_Step(){
  // update zeta
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    // update lower bound before updating new values
    for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
      lowerbound = lowerbound + (beta[v]-1) * (cache_digamma_var_zeta[k][v] - cache_digamma_var_zeta_k[k]);
      lowerbound = lowerbound - log_gamma(beta[v]);

      lowerbound = lowerbound - (var_zeta[k][v]-1) * (cache_digamma_var_zeta[k][v] - cache_digamma_var_zeta_k[k]);
      lowerbound = lowerbound + log_gamma(var_zeta[k][v]);
    }
    lowerbound = lowerbound + log_gamma(BetaSum);
    lowerbound = lowerbound - log_gamma(var_zeta_k[k]);

    var_zeta_k[k] = 0.0;
    for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
      var_zeta[k][v] = beta[v] + ss_zeta[k][v];
      var_zeta_k[k] = var_zeta_k[k] + var_zeta[k][v];
      
      ss_zeta[k][v] = 0.0;
      cache_digamma_var_zeta[k][v] = digamma(var_zeta[k][v]);
    }
    cache_digamma_var_zeta_k[k] = digamma(var_zeta_k[k]);
  }

  // update xi
  for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    // update lower bound
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      lowerbound = lowerbound + (alpha[k]-1) * (cache_digamma_var_xi[r][k] - cache_digamma_var_xi_r[r]);
      lowerbound = lowerbound - log_gamma(alpha[k]);

      lowerbound = lowerbound - (var_xi[r][k]-1) * (cache_digamma_var_xi[r][k] - cache_digamma_var_xi_r[r]);
      lowerbound = lowerbound + log_gamma(var_xi[r][k]);
    }
    lowerbound = lowerbound + log_gamma(AlphaSum);
    lowerbound = lowerbound - log_gamma(var_xi_r[r]);

    var_xi_r[r] = 0.0;
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      var_xi[r][k] = alpha[k] + ss_xi[r][k];
      var_xi_r[r] = var_xi_r[r] + var_xi[r][k];

      ss_xi[r][k] = 0.0;
      cache_digamma_var_xi[r][k] = digamma(var_xi[r][k]);
    }
    cache_digamma_var_xi_r[r] = digamma(var_xi_r[r]);
  }
  return 0;
}

void GeoSimple1Batch::Save_Expectation_Phi_VarZeta(){
  for(int t = 0; t < PuzaDEF::Instance()->TERM_NUM; t++){
    file_output << t << "\t";
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      file_output << k << ":" << (var_zeta[k][t] / var_zeta_k[k]) << " ";
    }
    file_output << endl;
  }
}

void GeoSimple1Batch::Save_Expectation_Pi_VarXi(){
  for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    file_output << geo_mapping[r] << "\t";
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      file_output << k << ":" << (var_xi[r][k] / var_xi_r[r]) << " ";
    }
    file_output << endl;
  }
}

void GeoSimple1Batch::Inference(){
  PuzaLogger::Instance()->PutString("Start Batch Variational Inference for Geo Simple Model I");
  // variational EM                                                                             
  var_iter = 0;
  double old_lowerbound = 0;
  while(1){
    lowerbound = 0.0;
    // E-step
    for(int d=0; d < PuzaDEF::Instance()->DOC_NUM; d++){
      //cout << d << endl;
      Single_E_Step(d);
    }
    //cout << lowerbound << endl;
    // M-step                                                                                                                                                                  
    Single_M_Step();

    if(var_iter % 20 == 0){
      string phi_file_name = PuzaDEF::Instance()->output_file_name + ".phi";
      file_output.open(phi_file_name.c_str());
      Save_Expectation_Phi_VarZeta();
      file_output.close();
      string pi_file_name = PuzaDEF::Instance()->output_file_name + ".pi";
      file_output.open(pi_file_name.c_str());
      Save_Expectation_Pi_VarXi();
      file_output.close();
    }

    PuzaLogger::Instance()->PutString("iter :" + to_string<int>(var_iter,std::dec) +  " lowerbound:" + to_string<double>(lowerbound,std::dec));
    
    if(var_iter > 100){
      if((lowerbound - old_lowerbound)<=0.00001)
        break;
    }
    if(var_iter >= 500)
      break;

    old_lowerbound = lowerbound;
    var_iter ++;
  }
  PuzaLogger::Instance()->PutString("Batch Learning Finished.");
}

void GeoSimple1Batch::Prediction(){
}
 
double GeoSimple1Batch::Likelihood(){
  return 0.0;
}

int GeoSimple1Batch::ReadDOCS(){
  int MAX_INT = 0;
  string temp;
  if(PuzaDEF::Instance()->input_file_name.empty()){
    cout << "Please give input file name." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->InputFile.open(PuzaDEF::Instance()->input_file_name.c_str());
  if(!PuzaDEF::Instance()->InputFile){
    cout << "Can't open the file." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->data_ids.clear();
  string original_id;
  string temp_geo_stamp;
  string temp_user_stamp;
  vector<string> buffer;
  while(!PuzaDEF::Instance()->InputFile.eof()){
    getline(PuzaDEF::Instance()->InputFile,temp,'\n');
    std::istringstream iss(temp);
    // get the doc id                                                                                                                                                 
    getline(iss,temp,'\t');
    original_id = temp;
    // get geo                                                                                                                                 
    getline(iss,temp,'\t');
    temp_geo_stamp = temp;
    // get user                                                                                                                                                  
    getline(iss,temp,'\t');
    temp_user_stamp = temp;
    // get the doc body
    getline(iss,temp,'\n');
    std::istringstream tempiss(temp);
    vector<string> tokens;
    copy(istream_iterator<string>(tempiss), istream_iterator<string>(), back_inserter<vector<string> >(tokens));
    if(tokens.size()<1){
      continue;
    }

    int now_id = PuzaDEF::Instance()->data_ids.size();
    PuzaDEF::Instance()->data_ids[now_id] = original_id;
    boost::unordered_map<string,int>::iterator geo_iter = geo_reverse_mapping.find(temp_geo_stamp);
    int geo_id = -1;
    int user_id = -1;
    if(geo_iter != geo_reverse_mapping.end()){
      // id found                                                                                                                                                                                                                       
      geo_id = (*geo_iter).second;
      doc_geo_stamp.push_back(geo_id);
    }
    else{
      geo_id = geo_reverse_mapping.size();
      geo_reverse_mapping[temp_geo_stamp] = geo_id;
      geo_mapping[geo_id] = temp_geo_stamp;
      doc_geo_stamp.push_back(geo_id);
    }

    boost::unordered_map<string,int>::iterator user_iter = user_reverse_mapping.find(temp_user_stamp);
    if(user_iter != user_reverse_mapping.end()){
      // id found                                                                                                                                                                                                                         
      user_id = (*user_iter).second;
      doc_user_stamp.push_back(user_id);
    }
    else{
      user_id = user_reverse_mapping.size();
      user_reverse_mapping[temp_user_stamp] = user_id;
      user_mapping[user_id] = temp_user_stamp;
      doc_user_stamp.push_back(user_id);
    }

    // check user geo dist
    if(user_id >= (int)var_lambda.size())
      var_lambda.resize(user_id+1);
    var_lambda[user_id][geo_id] = var_lambda[user_id][geo_id] + 1;
    TermList assignList;
    for (vector<string>::iterator iter = tokens.begin(); iter < tokens.end(); iter++){
      int word_id;
      int word_count;
      std::istringstream valueiss((*iter));
      getline(valueiss,temp,':');
      from_string<int>(word_id,temp,std::dec);
      getline(valueiss,temp);  // get count                                                                                                                                                                                                  
      from_string<int>(word_count,temp,std::dec);
      TermItem newAssign;
      newAssign.term_id = word_id;
      newAssign.term_stat = word_count;
      assignList.push_back(newAssign);
    }
    if((int)assignList.size() > MAX_INT)
      MAX_INT = assignList.size();
    TM_DOCS.push_back(assignList);
  }
  PuzaDEF::Instance()->InputFile.close();
  return MAX_INT;
}

int GeoSimple1Batch::GetDocuments(int choice){
  max_doc_len = ReadDOCS();
  allocate_memory<double>(var_nu, max_doc_len, PuzaDEF::Instance()->TOPIC_NUM);
  return 0;
}
#endif
